{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\pujing\\Anaconda3\\lib\\site-packages\\ipykernel\\parentpoller.py:116: UserWarning: Parent poll failed.  If the frontend dies,\n",
      "                the kernel may be left running.  Please let us know\n",
      "                about your system (bitness, Python, etc.) at\n",
      "                ipython-dev@scipy.org\n",
      "  ipython-dev@scipy.org\"\"\")\n"
     ]
    }
   ],
   "source": [
    "import torch.nn.functional as F\n",
    "import torch.nn as nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encode_conv_bn_x2(nn.Module):\n",
    "    def __init__(self, in_, out):\n",
    "        super(Encode_conv_bn_x2, self).__init__()\n",
    "        batchNorm_momentum = 0.1\n",
    "        self.relu    = nn.ReLU(inplace=True)\n",
    "        self.conv1 = nn.Conv2d(in_, out, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(out, momentum= batchNorm_momentum)\n",
    "        self.conv2 = nn.Conv2d(out, out, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(out, momentum= batchNorm_momentum)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = self.relu(x)\n",
    "        \n",
    "        x = self.conv2(x)\n",
    "        x = self.bn2(x)\n",
    "        x = self.relu(x)\n",
    "        return x    \n",
    "    \n",
    "class Encode_conv_bn_x3(nn.Module):\n",
    "    def __init__(self, in_, out):\n",
    "        super(Encode_conv_bn_x3, self).__init__()\n",
    "        batchNorm_momentum = 0.1\n",
    "        self.relu    = nn.ReLU(inplace=True)\n",
    "        self.conv1 = nn.Conv2d(in_, out, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(out, momentum= batchNorm_momentum)\n",
    "        self.conv2 = nn.Conv2d(out, out, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(out, momentum= batchNorm_momentum)\n",
    "        self.conv3 = nn.Conv2d(out, out, kernel_size=3, padding=1)\n",
    "        self.bn3 = nn.BatchNorm2d(out, momentum= batchNorm_momentum)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = self.relu(x)\n",
    "        \n",
    "        x = self.conv2(x)\n",
    "        x = self.bn2(x)\n",
    "        x = self.relu(x)\n",
    "        \n",
    "        x = self.conv3(x)\n",
    "        x = self.bn3(x)\n",
    "        x = self.relu(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dencode_conv_bn_x1(nn.Module):\n",
    "    def __init__(self, in_, out):\n",
    "        super(Dencode_conv_bn_x1, self).__init__()\n",
    "        batchNorm_momentum = 0.1\n",
    "        self.relu    = nn.ReLU(inplace=True)\n",
    "        self.conv1 = nn.Conv2d(in_, in_, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(in_, momentum= batchNorm_momentum)   \n",
    "        self.conv2= nn.Conv2d(in_, out, kernel_size=3, padding=1)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = self.relu(x)\n",
    "        \n",
    "        x = self.conv2(x)\n",
    "        return x\n",
    "    \n",
    "class Dencode_conv_bn_x2(nn.Module):\n",
    "    def __init__(self, in_, out):\n",
    "        super(Dencode_conv_bn_x2, self).__init__()\n",
    "        batchNorm_momentum = 0.1\n",
    "        \n",
    "        self.relu    = nn.ReLU(inplace=True)\n",
    "        self.conv1 = nn.Conv2d(in_, in_, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(in_, momentum= batchNorm_momentum)\n",
    "        \n",
    "        self.conv2= nn.Conv2d(in_, out, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(out, momentum= batchNorm_momentum)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = self.relu(x)\n",
    "        \n",
    "        x = self.conv2(x)\n",
    "        x = self.bn2(x)\n",
    "        x = self.relu(x)\n",
    "        return x\n",
    "\n",
    "class Dencode_conv_bn_x3(nn.Module):\n",
    "    def __init__(self, in_, out):\n",
    "        super(Dencode_conv_bn_x3, self).__init__()\n",
    "        batchNorm_momentum = 0.1\n",
    "        \n",
    "        self.relu    = nn.ReLU(inplace=True)\n",
    "        self.conv1 = nn.Conv2d(in_, in_, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(in_, momentum= batchNorm_momentum)\n",
    "        \n",
    "        self.conv2 = nn.Conv2d(in_, in_, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(in_, momentum= batchNorm_momentum)\n",
    "        \n",
    "        self.conv3 = nn.Conv2d(in_, out, kernel_size=3, padding=1)\n",
    "        self.bn3 = nn.BatchNorm2d(out, momentum= batchNorm_momentum)\n",
    "        \n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = self.bn2(x)\n",
    "        x = self.relu(x)\n",
    "        x = self.conv3(x)\n",
    "        x = self.bn3(x)\n",
    "        x = self.relu(x)\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "class SegNet(nn.Module):\n",
    "    def __init__(self,input_nbr,label_nbr):\n",
    "        super(SegNet, self).__init__()\n",
    "        batchNorm_momentum = 0.1\n",
    "        \n",
    "        self.encode1=Encode_conv_bn_x2(input_nbr,64)\n",
    "        self.encode2=Encode_conv_bn_x2(64,128)\n",
    "        self.encode3=Encode_conv_bn_x3(128,256)\n",
    "        self.encode4=Encode_conv_bn_x3(256,512)\n",
    "        self.encode5=Encode_conv_bn_x3(512,512)\n",
    "        \n",
    "        self.dencode5=Dencode_conv_bn_x3(512,512)\n",
    "        self.dencode4=Dencode_conv_bn_x3(512,256)\n",
    "        self.dencode3=Dencode_conv_bn_x2(256,128)\n",
    "        self.dencode2=Dencode_conv_bn_x2(128,64)\n",
    "        self.dencode1=Dencode_conv_bn_x1(64,label_nbr)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # Stage 1\n",
    "        x1=F.relu(self.encode1(x))\n",
    "        self.x1p, self.id1 = F.max_pool2d(x1,kernel_size=2, stride=2,return_indices=True)\n",
    "\n",
    "        # Stage 2\n",
    "        x2=F.relu(self.encode2(self.x1p))\n",
    "        self.x2p, self.id2 = F.max_pool2d(x2,kernel_size=2, stride=2,return_indices=True)\n",
    "\n",
    "        # Stage 3\n",
    "        x3=F.relu(self.encode3(self.x2p))\n",
    "        self.x3p, self.id3 = F.max_pool2d(x3,kernel_size=2, stride=2,return_indices=True)\n",
    "\n",
    "        # Stage 4\n",
    "        x4=F.relu(self.encode4(self.x3p))\n",
    "        self.x4p, self.id4 = F.max_pool2d(x4,kernel_size=2, stride=2,return_indices=True)\n",
    "\n",
    "        # Stage 5\n",
    "        x5=F.relu(self.encode5(self.x4p))\n",
    "        self.x5p, self.id5 = F.max_pool2d(x5,kernel_size=2, stride=2,return_indices=True)\n",
    "        \n",
    "        # Stage 5d\n",
    "        x5 = F.max_unpool2d(self.x5p, self.id5, kernel_size=2, stride=2)\n",
    "        x5=F.relu(self.dencode5(x5))\n",
    "\n",
    "        # Stage 4d\n",
    "        x4= F.max_unpool2d(x5, self.id4, kernel_size=2, stride=2)\n",
    "        x4=F.relu(self.dencode4(x4))\n",
    "        \n",
    "        \n",
    "        # Stage 3d\n",
    "        x3= F.max_unpool2d(x4, self.id3, kernel_size=2, stride=2)\n",
    "        x3=F.relu(self.dencode3(x3))\n",
    "\n",
    "        # Stage 2d\n",
    "        x2= F.max_unpool2d(x3, self.id2, kernel_size=2, stride=2)\n",
    "        x2=F.relu(self.dencode2(x2))\n",
    "\n",
    "        # Stage 1d\n",
    "        x1 = F.max_unpool2d(x2, self.id1, kernel_size=2, stride=2)\n",
    "        x1=self.dencode1(x1)\n",
    "        return x1\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------\n",
      "        Layer (type)               Output Shape         Param #\n",
      "================================================================\n",
      "            Conv2d-1         [-1, 64, 224, 224]           1,792\n",
      "       BatchNorm2d-2         [-1, 64, 224, 224]             128\n",
      "              ReLU-3         [-1, 64, 224, 224]               0\n",
      "            Conv2d-4         [-1, 64, 224, 224]          36,928\n",
      "       BatchNorm2d-5         [-1, 64, 224, 224]             128\n",
      "              ReLU-6         [-1, 64, 224, 224]               0\n",
      " Encode_conv_bn_x2-7         [-1, 64, 224, 224]               0\n",
      "            Conv2d-8        [-1, 128, 112, 112]          73,856\n",
      "       BatchNorm2d-9        [-1, 128, 112, 112]             256\n",
      "             ReLU-10        [-1, 128, 112, 112]               0\n",
      "           Conv2d-11        [-1, 128, 112, 112]         147,584\n",
      "      BatchNorm2d-12        [-1, 128, 112, 112]             256\n",
      "             ReLU-13        [-1, 128, 112, 112]               0\n",
      "Encode_conv_bn_x2-14        [-1, 128, 112, 112]               0\n",
      "           Conv2d-15          [-1, 256, 56, 56]         295,168\n",
      "      BatchNorm2d-16          [-1, 256, 56, 56]             512\n",
      "             ReLU-17          [-1, 256, 56, 56]               0\n",
      "           Conv2d-18          [-1, 256, 56, 56]         590,080\n",
      "      BatchNorm2d-19          [-1, 256, 56, 56]             512\n",
      "             ReLU-20          [-1, 256, 56, 56]               0\n",
      "           Conv2d-21          [-1, 256, 56, 56]         590,080\n",
      "      BatchNorm2d-22          [-1, 256, 56, 56]             512\n",
      "             ReLU-23          [-1, 256, 56, 56]               0\n",
      "Encode_conv_bn_x3-24          [-1, 256, 56, 56]               0\n",
      "           Conv2d-25          [-1, 512, 28, 28]       1,180,160\n",
      "      BatchNorm2d-26          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-27          [-1, 512, 28, 28]               0\n",
      "           Conv2d-28          [-1, 512, 28, 28]       2,359,808\n",
      "      BatchNorm2d-29          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-30          [-1, 512, 28, 28]               0\n",
      "           Conv2d-31          [-1, 512, 28, 28]       2,359,808\n",
      "      BatchNorm2d-32          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-33          [-1, 512, 28, 28]               0\n",
      "Encode_conv_bn_x3-34          [-1, 512, 28, 28]               0\n",
      "           Conv2d-35          [-1, 512, 14, 14]       2,359,808\n",
      "      BatchNorm2d-36          [-1, 512, 14, 14]           1,024\n",
      "             ReLU-37          [-1, 512, 14, 14]               0\n",
      "           Conv2d-38          [-1, 512, 14, 14]       2,359,808\n",
      "      BatchNorm2d-39          [-1, 512, 14, 14]           1,024\n",
      "             ReLU-40          [-1, 512, 14, 14]               0\n",
      "           Conv2d-41          [-1, 512, 14, 14]       2,359,808\n",
      "      BatchNorm2d-42          [-1, 512, 14, 14]           1,024\n",
      "             ReLU-43          [-1, 512, 14, 14]               0\n",
      "Encode_conv_bn_x3-44          [-1, 512, 14, 14]               0\n",
      "           Conv2d-45          [-1, 512, 14, 14]       2,359,808\n",
      "      BatchNorm2d-46          [-1, 512, 14, 14]           1,024\n",
      "             ReLU-47          [-1, 512, 14, 14]               0\n",
      "           Conv2d-48          [-1, 512, 14, 14]       2,359,808\n",
      "      BatchNorm2d-49          [-1, 512, 14, 14]           1,024\n",
      "             ReLU-50          [-1, 512, 14, 14]               0\n",
      "           Conv2d-51          [-1, 512, 14, 14]       2,359,808\n",
      "      BatchNorm2d-52          [-1, 512, 14, 14]           1,024\n",
      "             ReLU-53          [-1, 512, 14, 14]               0\n",
      "Dencode_conv_bn_x3-54          [-1, 512, 14, 14]               0\n",
      "           Conv2d-55          [-1, 512, 28, 28]       2,359,808\n",
      "      BatchNorm2d-56          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-57          [-1, 512, 28, 28]               0\n",
      "           Conv2d-58          [-1, 512, 28, 28]       2,359,808\n",
      "      BatchNorm2d-59          [-1, 512, 28, 28]           1,024\n",
      "             ReLU-60          [-1, 512, 28, 28]               0\n",
      "           Conv2d-61          [-1, 256, 28, 28]       1,179,904\n",
      "      BatchNorm2d-62          [-1, 256, 28, 28]             512\n",
      "             ReLU-63          [-1, 256, 28, 28]               0\n",
      "Dencode_conv_bn_x3-64          [-1, 256, 28, 28]               0\n",
      "           Conv2d-65          [-1, 256, 56, 56]         590,080\n",
      "      BatchNorm2d-66          [-1, 256, 56, 56]             512\n",
      "             ReLU-67          [-1, 256, 56, 56]               0\n",
      "           Conv2d-68          [-1, 128, 56, 56]         295,040\n",
      "      BatchNorm2d-69          [-1, 128, 56, 56]             256\n",
      "             ReLU-70          [-1, 128, 56, 56]               0\n",
      "Dencode_conv_bn_x2-71          [-1, 128, 56, 56]               0\n",
      "           Conv2d-72        [-1, 128, 112, 112]         147,584\n",
      "      BatchNorm2d-73        [-1, 128, 112, 112]             256\n",
      "             ReLU-74        [-1, 128, 112, 112]               0\n",
      "           Conv2d-75         [-1, 64, 112, 112]          73,792\n",
      "      BatchNorm2d-76         [-1, 64, 112, 112]             128\n",
      "             ReLU-77         [-1, 64, 112, 112]               0\n",
      "Dencode_conv_bn_x2-78         [-1, 64, 112, 112]               0\n",
      "           Conv2d-79         [-1, 64, 224, 224]          36,928\n",
      "      BatchNorm2d-80         [-1, 64, 224, 224]             128\n",
      "             ReLU-81         [-1, 64, 224, 224]               0\n",
      "           Conv2d-82          [-1, 4, 224, 224]           2,308\n",
      "Dencode_conv_bn_x1-83          [-1, 4, 224, 224]               0\n",
      "================================================================\n",
      "Total params: 28,854,724\n",
      "Trainable params: 28,854,724\n",
      "Non-trainable params: 0\n",
      "----------------------------------------------------------------\n",
      "Input size (MB): 0.57\n",
      "Forward/backward pass size (MB): 557.38\n",
      "Params size (MB): 110.07\n",
      "Estimated Total Size (MB): 668.02\n",
      "----------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "seg=SegNet(3,4)\n",
    "from torchsummary import summary\n",
    "summary(seg, (3,224,224))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
